Skip to content

Add Unsloth Docker images (base + Studio) for any NVIDIA GPU host, Ampere through Blackwell#5748

Open
danielhanchen wants to merge 120 commits into
mainfrom
docker-blackwell-build
Open

Add Unsloth Docker images (base + Studio) for any NVIDIA GPU host, Ampere through Blackwell#5748
danielhanchen wants to merge 120 commits into
mainfrom
docker-blackwell-build

Conversation

@danielhanchen

@danielhanchen danielhanchen commented May 24, 2026

Copy link
Copy Markdown
Member

Summary

Adds a Docker setup for Unsloth that runs on any NVIDIA GPU host from Ampere through Blackwell (sm_80 through sm_120: A100, RTX 30/40, H100, B100/B200, RTX 50-series, RTX 6000 Pro Blackwell) and natively on aarch64 (GB10 / Grace, DGX Spark). Two images are published to docker.io/unsloth/unsloth:

  • Lean base image (docker/Dockerfile, tag :base): the full training stack -- torch 2.10.0+cu128, Unsloth + unsloth_zoo, TRL / PEFT / accelerate, bitsandbytes, triton, xformers (amd64), vLLM -- plus llama.cpp for GGUF tooling, JupyterLab, and the baked Unsloth notebooks. Run it headless for training, unsloth-run <notebook|url>, or jupyter lab.
  • Full image (docker/Dockerfile.studio, tag :latest): layers Unsloth Studio on top of the base image and runs the production service trio under supervisord -- Studio on 8000, JupyterLab on 8888, key-only sshd on 22 -- plus an optional Cloudflare tunnel for JupyterLab.

The build itself requires no GPU at all: it runs on free GitHub-hosted runners, on a developer laptop without an NVIDIA card, or on any datacenter GPU. All produce byte-identical images.

Multi-arch. amd64 and arm64 are built in parallel on native GitHub runners (ubuntu-latest + ubuntu-24.04-arm, both free on public repos since Aug 2025), pushed by digest, then merged into a single multi-platform manifest, so docker pull selects the right child automatically. Native arm64 is ~3x faster than QEMU and runs on DGX Spark / Grace with CUDA working as normal (no runtime emulation).

Why the build does not need a GPU

There are four places where a naive Docker build silently couples to the build-host GPU. The Dockerfile breaks each one:

1. Wheel selection. The README install line uv pip install unsloth --torch-backend=auto introspects the build host's driver. This Dockerfile pins torch==2.10.0 against --extra-index-url https://download.pytorch.org/whl/cu128 explicitly. No --torch-backend=auto, no install.sh.

2. Dep resolution order. Splitting installs into multiple pip install calls lets bitsandbytes 0.49.x's transitive cuda-toolkit==13.0.2 dep silently upgrade torch 2.10.0+cu128 -> 2.12.0+cu130 in a later pass, leaving cu128 xformers stranded. This Dockerfile collapses everything into a single uv pip install with --index-strategy unsafe-best-match so the resolver sees all constraints at once.

3. Build-time verification. torch.cuda.get_arch_list() returns [] when no GPU is visible. The Dockerfile uses the raw C++ accessor torch._C._cuda_getArchFlags(), which reads compiled wheel metadata directly ('sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120' on amd64; 'sm_80 sm_90 sm_100 sm_120' on aarch64). Required packages are checked via importlib.metadata.version() instead of importing them, because import unsloth triggers torch.cuda.get_device_properties(0), which can't be satisfied on a GPU-less host. Import-time correctness is exercised at deploy time by smoke_test.py.

4. Compiled-kernel cache. If anything imports unsloth during the build, Triton JITs kernels keyed to the build host's compute capability and bakes them into unsloth_compiled_cache/. UNSLOTH_COMPILE_DISABLE=1 and UNSLOTH_COMPILE_OVERWRITE=0 prevent this. The deploy GPU produces its own cache on first use.

The underlying reason all of this works is that cu128 PyTorch wheels are already fat binaries, cross-compiled upstream to include SASS for every architecture from sm_70 through sm_120 (and sm_80/90/100/120 on aarch64). The build host's GPU was never needed for the binary content, only by code that pretended to need it.

What is in the images

Base image (:base)

  • Training stack pinned at torch 2.10.0+cu128 (held against the cu cascade), with TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;10.3;12.0+PTX" so any source build inside the container honours the full Ampere-through-Blackwell range. The +PTX suffix on the top arch gives forward-compat JIT-PTX for future consumer Blackwell SKUs.
  • vLLM and llama.cpp (prebuilt) for inference and GGUF tooling.
  • The curated notebook extras (JupyterLab, ipywidgets, audio/TTS codecs, etc.) pinned to tested versions for reproducible rebuilds.
  • unslothai/notebooks baked in as a read-only template and synced to /workspace/unsloth-notebooks on boot (edit-preserving refresh from GitHub when reachable).
  • Notebook ergonomics: a !pip/!uv shim that keeps the core GPU stack pinned while letting a notebook install its own extras, per-notebook transformers-version sidecars, and unsloth-run for headless execution of any notebook or URL.
  • arm64 swaps to the [huggingface] extra (no cu128 aarch64 xformers/vLLM wheel) and installs cuda-nvrtc/nvcc from NVIDIA's sbsa repo.

Full image (:latest)

  • Everything in the base image, plus Unsloth Studio.
  • supervisord runs Studio (8000), JupyterLab (8888), an optional Cloudflare tunnel for JupyterLab, and key-only sshd (22). The studio venv dedups its CUDA libraries against the base venv to keep the image size down.

arm64 / aarch64

A native arm64 child is published in the same manifest. CI covers it two ways: docker-build-arm64-native.yml builds + smoke-checks on the free ubuntu-24.04-arm runner (the path DGX Spark / Grace users take), and docker-build-arm64-qemu.yml cross-builds under QEMU as a fallback signal and to validate the documented setup_qemu.sh recipe.

Validation

Base image validated end-to-end on a B200 (sm_100) host. With CUDA_VISIBLE_DEVICES="" (simulating the GPU-less CI runner), the resolved stack is torch 2.10.0+cu128, triton, xformers (cu128 wheel from the PyTorch index), bitsandbytes, unsloth + unsloth_zoo, transformers, trl, peft, accelerate, with arch flags ['sm_70','sm_75','sm_80','sm_86','sm_90','sm_100','sm_120']. Runtime path on B200 GPU 0: smoke_test.py imports unsloth, loads Llama-3.2-1B-Instruct-bnb-4bit in 4-bit, and completes LoRA steps with loss decreasing. The full image was launched on the same host and Studio, JupyterLab, and sshd all came up under supervisord.

Files

Path Purpose
docker/Dockerfile multi-stage cu128 base image, no GPU required (amd64 + arm64)
docker/Dockerfile.studio full image: base + Studio + JupyterLab + sshd under supervisord
docker/entrypoint.sh, docker/supervisord.conf, docker/studio_launch.sh boot + service orchestration
docker/unsloth_sync_notebooks.sh, docker/unsloth_nb_compat.py, docker/unsloth_nb_content_sig.py baked-notebook sync + per-notebook transformers sidecars
docker/unsloth_pip_shim.py, docker/unsloth_run.py, docker/unsloth_ipython_startup.py !pip/!uv shim, headless runner, IPython startup hook
docker/unsloth_studio_update.sh, docker/unsloth_llama_update.sh, docker/unsloth_jupyter_tunnel.sh, docker/fetch_llama_prebuilt.py in-image updaters + llama.cpp fetch + Cloudflare tunnel
docker/run.sh, docker/build.sh, docker/test_locally.sh, docker/docker_confirm.sh, docker/docker_confirm.ps1, docker/setup_qemu.sh local build / run / cross-OS confirm helpers
docker/freeze.sh, docker/hf_pull.sh, docker/hf_push.sh, docker/smoke_test.py, docker/.dockerignore lockfile extract, HF helpers, runtime smoke test
.github/workflows/docker-publish.yml multi-arch build + push of both images to Docker Hub

Design notes

  • Base images: nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 for the build stage, -cudnn-runtime-ubuntu24.04 for deploy. No nvcc in the published amd64 image.
  • A lockfile is emitted at /opt/unsloth-venv/requirements.lock.txt inside the image and can be extracted with docker/freeze.sh for a fully-pinned rebuild later.

Test plan

  • Set repo secrets DOCKERHUB_USERNAME and DOCKERHUB_TOKEN
  • Trigger the publish workflow (or merge to main) and confirm both images build on the native amd64 + arm64 runners in under ~90 min and merge into one manifest
  • Pull unsloth/unsloth:latest on an RTX 50-series host and confirm Studio (8000) + JupyterLab (8888) come up
  • Pull the same tag on a B200 host and confirm the image works with no rebuild
  • Pull on a DGX Spark / Grace (arm64) host and confirm the arm64 child runs natively
  • Optional: set vars.HAS_GPU_RUNNER=true once a self-hosted GPU runner is registered so the post-publish smoke test exercises real sm_120 paths

danielhanchen and others added 2 commits May 24, 2026 06:52
Adds a multi-stage Dockerfile producing an image that works on Ampere through
Blackwell (sm_80 through sm_120: A100, RTX 30/40, H100, B100/B200, RTX 50-series,
RTX 6000 Pro Blackwell). The build itself requires no GPU at all and runs on a
free GitHub-hosted ubuntu-latest runner.

How the GPU-less build works:

1. cu128 PyTorch wheels are fat binaries. torch._C._cuda_getArchFlags() returns
   'sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120' regardless of which GPU
   compiled the image, because the wheels are cross-compiled upstream by the
   PyTorch team.

2. All deps resolve in a single uv pip install pass with explicit pins
   (torch==2.10.0, --extra-index-url cu128, no --torch-backend=auto, no
   install.sh). This prevents the silent cu cascade where bitsandbytes'
   transitive cuda-toolkit==13 dep upgrades torch to 2.12+cu130 in a later
   resolver pass, leaving xformers and other cu128 wheels stranded.

3. Build-time verification uses package metadata (importlib.metadata.version)
   and the raw torch._C._cuda_getArchFlags() accessor. We deliberately avoid
   import unsloth at build time because unsloth.__init__ calls
   torch.cuda.get_device_properties(0), which requires an actual CUDA device
   and is not bypassable. Import-time correctness is exercised at deploy time
   by smoke_test.py with --gpus all.

4. UNSLOTH_COMPILE_DISABLE=1 and CUDA_VISIBLE_DEVICES="" during the build stage
   prevent any code path from JIT-compiling kernels for the build host's
   compute capability and baking the resulting cache into the image. The
   deploy GPU produces its own cache on first use.

Other notes:

- --index-strategy unsafe-best-match is needed because the PyTorch wheel index
  serves an old requests==2.28.1 that conflicts with datasets>=2.32.2, which
  the default first-index-wins strategy rejects.
- Extra is cu128-ampere-torch2100 (ampere precedes the torch version in the
  pyproject ordering).
- No flash-attn in the base image. FA3 is hard-refused on Blackwell upstream
  and unsloth gracefully falls back to xformers + SDPA. Users on Ampere /
  Ada / Hopper who want FA2 can pip install flash-attn on top.
- Two stages: nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 for the build,
  -cudnn-runtime for the deploy image. No nvcc in the published image.
- A lockfile is emitted at /opt/unsloth-venv/requirements.lock.txt inside
  the image and can be extracted with docker/freeze.sh for byte-identical
  rebuilds even after PyPI moves on.

CI workflow .github/workflows/docker-publish.yml:

- Builds on ubuntu-latest on every push to main, every tag, weekly via cron,
  and manually via workflow_dispatch. Pushes to docker.io/unsloth/unsloth
  with cache via type=gha.
- Optional smoke-test job runs on a self-hosted GPU runner if vars.HAS_GPU_RUNNER
  is set; skipped otherwise. End-to-end verification on sm_120 hardware is a
  nice-to-have, not a publish blocker.

Validation:

- Install path validated on a B200 host with CUDA_VISIBLE_DEVICES="" set
  (simulating the GPU-less CI runner): torch 2.10.0+cu128 holds, xformers
  0.0.34, bitsandbytes 0.49.2, triton 3.6.0, transformers 5.5.0, trl 0.24.0,
  peft 0.19.1, accelerate 1.13.0. Arch flags include sm_100 and sm_120.
- Runtime path validated end-to-end on B200: smoke_test.py imports unsloth,
  loads Llama-3.2-1B-Instruct-bnb-4bit in 4-bit, completes 5 LoRA steps with
  loss decreasing 4.11 -> 3.75. xformers fallback active as designed.

Files:

- docker/Dockerfile             multi-stage cu128 build
- docker/build.sh               local build wrapper
- docker/freeze.sh              extract lockfile from a built image
- docker/smoke_test.py          runtime verification, run with --gpus all
- docker/.dockerignore
- .github/workflows/docker-publish.yml
Comment thread .github/workflows/docker-publish.yml Fixed

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c6d92160f6

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread .github/workflows/docker-publish.yml Outdated
Comment on lines +115 to +117
docker pull ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest
docker run --rm --gpus all \
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest \

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Smoke-test the image built in this run, not latest

The smoke-test job always pulls :latest, but this workflow also runs on tag pushes where latest is not guaranteed to be among the tags produced by metadata-action (it is only enabled on the default branch in this workflow). In that case, the smoke test can validate an older image and miss regressions in the freshly built tag from the current run.

Useful? React with 👍 / 👎.

Comment thread .github/workflows/docker-publish.yml Outdated
Comment on lines +96 to +97
UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || 'main' }}
UNSLOTH_ZOO_REF=${{ github.event.inputs.unsloth_zoo_ref || 'main' }}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Build from triggering ref instead of hardcoding main

For non-workflow_dispatch events (including tag pushes), github.event.inputs.* is unset, so these build args always resolve to main. That means images produced for v* tags can contain unsloth and unsloth-zoo code from main rather than the release ref that triggered the run, which breaks release correctness and reproducibility.

Useful? React with 👍 / 👎.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a Dockerized environment for Unsloth and unsloth-zoo, specifically optimized for NVIDIA Blackwell GPUs (sm_100 and sm_120). The changes include a multi-stage Dockerfile, build and freeze scripts, and a comprehensive smoke test to verify GPU compatibility and training functionality. Review feedback suggests optimizing the Dockerfile by removing a redundant installation of the uv tool, correcting a version mismatch for torchaudio to ensure consistency with the PyTorch stack, and relocating cache directories outside of the workspace to prevent issues when mounting host volumes at runtime.

Comment thread docker/Dockerfile Outdated
Comment on lines +99 to +101
RUN ${VENV}/bin/pip install uv \
&& ${VENV}/bin/uv pip install \
--python ${VENV}/bin/python \

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The uv tool is already installed at the system level in line 63. Installing it again inside the virtual environment at line 99 is redundant. Using the system-wide uv binary to install packages into the venv is more efficient and avoids unnecessary layers.

RUN uv pip install \
        --python ${VENV}/bin/python \

Comment thread docker/Dockerfile Outdated
--python ${VENV}/bin/python \
--index-strategy unsafe-best-match \
--extra-index-url https://download.pytorch.org/whl/cu128 \
"torch==2.10.0" "torchvision==0.25.0" "torchaudio==2.11.0" \

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

There appears to be a version mismatch for torchaudio. While torch is pinned to 2.10.0 and torchvision to 0.25.0 (which correctly follows the standard 0.(Y+15) mapping for Torch 2.10), torchaudio is set to 2.11.0. Typically, PyTorch and Torchaudio versions are released in sync (e.g., Torch 2.6.0 with Torchaudio 2.6.0). Using 2.10.0 ensures consistency across the stack.

        "torch==2.10.0" "torchvision==0.25.0" "torchaudio==2.10.0" \

Comment thread docker/Dockerfile
Comment on lines +177 to +178
HF_HOME=/workspace/.cache/huggingface \
TRITON_CACHE_DIR=/workspace/.cache/triton \

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Setting HF_HOME and TRITON_CACHE_DIR to subdirectories of /workspace (the WORKDIR) can lead to issues when users mount a host directory to /workspace at runtime. The mount will obscure the directories created during the build, forcing the application to recreate them at runtime, which can cause permission issues or redundant downloads. Moving these caches to a location outside of the workspace, such as /opt/cache, avoids these issues.

    HF_HOME=/opt/cache/huggingface \
    TRITON_CACHE_DIR=/opt/cache/triton \

Comment thread docker/Dockerfile
COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv

WORKDIR /workspace
RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

When moving the cache directories outside of /workspace, creating them with broad permissions (e.g., 777) ensures that non-root users can write to the cache at runtime without encountering permission errors.

RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR} && chmod -R 777 /opt/cache

When someone launches the unsloth container, the common failure modes are not
unsloth bugs -- they're Docker / nvidia-container-toolkit / driver issues that
surface as cryptic CUDA errors deep in torch. The entrypoint catches the three
that cover ~95% of "it doesn't work" reports up front:

1. nvidia-smi inside the container sees no GPU
   -> user forgot --gpus all, or host is missing nvidia-container-toolkit
   -> entrypoint prints the exact docker run flag and the toolkit install URL
2. nvidia-smi works but torch.cuda.is_available() is False
   -> host driver is older than CUDA 12.8 supports
   -> entrypoint prints the minimum driver version per architecture
3. compute capability < sm_80
   -> entrypoint prints the supported architecture table and exits

Each check fails with a clear, actionable message rather than a stack trace.
Set UNSLOTH_SKIP_GPU_CHECK=1 to bypass (for docs builds, offline tooling, CI).

run.sh wraps `docker run` with the flags people most often forget:
  --gpus all           (without it, the new entrypoint refuses to start)
  --ipc=host           (DataLoader workers need >64MB shm)
  --ulimit memlock=-1  (NCCL + CUDA pinned host buffers)
  --ulimit stack=64MB  (some torch kernels OOM the default 8MB stack)

Plus it mounts the host HF cache + Triton JIT cache so model downloads and
compiled kernels persist across container runs, and forwards HF_TOKEN /
WANDB_API_KEY / UNSLOTH_LICENSE only when they are set on the host.

Usage:
  bash docker/run.sh                                  # interactive python REPL
  bash docker/run.sh bash                             # shell in container
  bash docker/run.sh python /workspace/smoke_test.py
  bash docker/run.sh python /workspace/host/train.py  # $PWD mounted at /workspace/host

Verified locally:
- No GPU visible: entrypoint refuses with driver-version message, exit 1
- B200 sm_100 visible: entrypoint prints GPU banner, exits cleanly into the
  user command (rc=0)

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 58693c4c73

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread docker/run.sh Outdated
Comment on lines +55 to +59
[[ -n "${HF_TOKEN:-}" ]] && ENV_FORWARD+=(-e "HF_TOKEN=${HF_TOKEN}")
[[ -n "${WANDB_API_KEY:-}" ]] && ENV_FORWARD+=(-e "WANDB_API_KEY=${WANDB_API_KEY}")
[[ -n "${UNSLOTH_LICENSE:-}" ]] && ENV_FORWARD+=(-e "UNSLOTH_LICENSE=${UNSLOTH_LICENSE}")

set -x

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Remove xtrace before invoking docker run with secrets

This script conditionally forwards HF_TOKEN, WANDB_API_KEY, and UNSLOTH_LICENSE, then enables set -x right before docker run, which prints the fully expanded command line. In any environment where those variables are set (local terminals with history/log capture or CI logs), their raw values are exposed in plaintext, creating an avoidable credential leak.

Useful? React with 👍 / 👎.

Comment thread docker/smoke_test.py
Comment on lines +61 to +63
import transformers

print(f"transformers {transformers.__version__}")

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Import unsloth before transformers in smoke test

The default execution path runs training (check_tiny_train) but check_imports imports transformers first, which contradicts the later requirement that Unsloth be imported first for patched training behavior. This means the smoke test is not validating the intended Unsloth training path and can produce misleading pass/fail results (e.g., patches not applied or warnings/fallback paths during the LoRA step).

Useful? React with 👍 / 👎.

Single bash script that runs the full validation flow against the image:

  1. Host pre-flight: docker version, nvidia-smi, nvidia-container-toolkit
     runtime registered with docker.
  2. Build the image (auto-detects the build context -- current dir,
     docker/ subdir, or clones the docker-blackwell-build branch into
     /tmp/unsloth-pr/).
  3a. Smoke test: 5-step LoRA on Llama-3.2-1B-Instruct-bnb-4bit.
  3b. Real workload: gpt-oss-20B fine-tuning notebook from
      unslothai/notebooks, patched to max_steps=10, with the three
      pre-train demo generations dropped for brevity. Auto-installs
      triton_kernels at the SHA the upstream notebook pins for MXFP4.

All output is teed to /tmp/unsloth-docker-test/ (or --log-dir).

Usage:
  bash docker/test_locally.sh                  # full run, ~15 min
  bash docker/test_locally.sh --skip-notebook  # blocks 1-3a only, ~3 min
  bash docker/test_locally.sh --skip-build     # reuse existing TAG
  TAG=my:tag HF_TOKEN=hf_xxx bash docker/test_locally.sh

Each block fails fast with the exact log path to paste back.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: acbb16c8a1

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread docker/Dockerfile Outdated
--python ${VENV}/bin/python \
--index-strategy unsafe-best-match \
--extra-index-url https://download.pytorch.org/whl/cu128 \
"torch==2.10.0" "torchvision==0.25.0" "torchaudio==2.11.0" \

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Align torchaudio pin with torch 2.10 stack

The install set mixes torch==2.10.0 with torchaudio==2.11.0, but TorchAudio wheels are release-coupled to specific Torch versions; this can make dependency resolution fail or force an unexpected Torch change, which breaks the Docker build’s stated guarantee that Torch stays on 2.10.0. This is especially risky here because the same layer also installs many transitive deps from multiple indexes, so one incompatible pin can fail the image build on CI.

Useful? React with 👍 / 👎.

Comment thread docker/entrypoint.sh Outdated
Comment on lines +28 to +31
if ! command -v nvidia-smi >/dev/null 2>&1; then
err "nvidia-smi not found inside the container."
err "The CUDA runtime in this image is broken. Re-pull the image."
exit 1

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Avoid hard-failing when nvidia-smi binary is unavailable

Startup currently exits before any CUDA check if nvidia-smi is missing, but there are valid GPU runtimes (for example compute-only capability profiles) where CUDA is usable while nvidia-smi/NVML tools are not mounted. In those environments the container will refuse to start even though torch.cuda.is_available() could succeed, causing false negatives in production deployments that intentionally limit driver capabilities.

Useful? React with 👍 / 👎.

The Dockerfile uses BuildKit-only features (the # syntax=docker/dockerfile:1.7
parser directive and RUN ... <<'PY' heredocs added in dockerfile 1.3+). The
legacy builder rejects the --progress flag at the CLI level and would fail
later at the heredocs anyway.

Detect docker buildx and use it when available (preserves --progress=plain
output). Otherwise fall back to plain `docker build` with DOCKER_BUILDKIT=1
exported, which gets the BuildKit features without buildx's nicer formatting.

Reproduces the failure path seen on Docker 28.2.2 without buildx installed:
  unknown flag: --progress
  ERROR docker build exited 125
Docker 28 removed the legacy image builder entirely. Setting
DOCKER_BUILDKIT=1 no longer falls back to a builtin builder -- it
delegates to buildx, which then errors out if buildx isn't installed:

  ERROR: BuildKit is enabled but the buildx component is missing
         or broken.

The Ubuntu docker.io package omits buildx by default, so users on
that path hit this immediately. Detect missing buildx up front and
print exact install commands for apt / dnf / manual binary instead
of attempting a fallback that cannot work.
If the user is not in the 'docker' group, every docker command after the
pre-flight returns "permission denied while trying to connect to the Docker
daemon socket at /var/run/docker.sock". This used to surface as a confusing
buildx failure mid-Block-2, but the actual problem is a host permissions
issue that's settable up front.

Detect by running 'docker info' and checking its exit code (not just grep
on its output -- a permission failure prints to stderr and returns non-zero,
so the old grep-based check was a silent skip).

Also clarify the nvidia-runtime WARN: on Docker 28+ with CDI mode this is
a false positive most of the time. The real GPU-attach test is the smoke
run in Block 3a, where the container entrypoint catches missing GPUs with
an actionable message.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 23a5b43180

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread docker/test_locally.sh
cd /workspace/host

echo "=== install triton_kernels (MXFP4 support for unsloth/gpt-oss-20b) ==="
pip install -q 'git+https://github.com/triton-lang/triton.git@0add68262ab0a2e33b84524346cb27cbb2787356#subdirectory=python/triton_kernels' 2>&1 | tail -5

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Propagate pip install failures through the tail pipeline

The generated run_notebook.sh script runs with set -e but not pipefail, and this line pipes pip install into tail -5; in Bash that means the pipeline exits with tail's status, so a failed pip install can be treated as success and the script continues until later steps fail with misleading errors (for example missing triton_kernels imports). This causes false-positive notebook validation and makes debugging CI/local failures much harder whenever installation fails (network hiccups, dependency conflicts, or bad commit hash).

Useful? React with 👍 / 👎.

Ubuntu 24.04 (noble) marks the system Python interpreter as
externally-managed per PEP 668, so:

  curl get-pip.py | python
  python -m pip install -U pip uv

fails inside the builder image with:

  error: externally-managed-environment
  This environment is externally managed

The system-level pip and uv were never used: the very next RUN creates
the venv at /opt/unsloth-venv, which bootstraps its own pip via the
ensurepip module (provided by the python3.12-venv apt package). uv is
then installed INTO the venv with the venv's pip, and used from there.

Drop the two system-pip bootstrap lines. The venv path is unchanged.

Reproduces on any Docker build of the unsloth-blackwell image against
a noble base image (which our nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
is).
… / peft

unsloth_zoo/__init__.py guards against being imported standalone:

  if "UNSLOTH_IS_PRESENT" not in os.environ:
      raise ImportError("Please install Unsloth via `pip install unsloth`!")

The env var is set by unsloth/__init__.py at import time, so importing
unsloth must happen first. The old check_imports() imported xformers,
bnb, transformers, trl, peft, then unsloth_zoo -- which fired the guard
because unsloth had not been imported yet.

Reorder check_imports() to import unsloth (and unsloth_zoo) first, then
the rest. check_unsloth_import() becomes a thin re-import to keep the
"FastLanguageModel reachable" banner in the output.

Same fix the unsloth README has been recommending for years: "import
unsloth at the top of your file, before transformers/trl/peft."
Triton's nvidia backend lazily JIT-compiles a small C extension
(CudaUtils, in triton/backends/nvidia/driver.py) on first GPU access.
Without a C compiler and Python headers in the runtime image, the
very first forward pass of any Unsloth model dies with:

  RuntimeError: Failed to find C compiler.
                Please specify via CC environment variable.

The builder stage has build-essential and python3.12-dev so this
worked during the build's verification step (no GPU = no Triton kernel
call = no C extension build). But the runtime stage stripped those
out for size, so the failure only surfaces when a real user runs
training inside the container.

Add gcc + g++ + python3.12-dev to the runtime stage. Increases the
runtime image by ~250MB, which is the cost of letting Triton JIT
correctly. Pre-compiling CudaUtils at build time would need a real
CUDA device (the constructor calls cuda runtime functions), so
shipping the toolchain is the right trade-off.
@danielhanchen

Copy link
Copy Markdown
Member Author

Smoke-test validation on a fresh deploy host (AWS B200, not the build host)

End-to-end validated docker/test_locally.sh --skip-notebook on an AWS EC2 instance with 8x B200, driver 590.48.01, Docker 28.2.2 (Ubuntu docker.io + docker-buildx 0.30.1). The build host was a separate GCP B200 — so this confirms cross-host reproducibility.

What was validated

  • Build: builder + runtime stages completed cleanly
  • Built-in arch check: arches: ['sm_70', 'sm_75', 'sm_80', 'sm_86', 'sm_90', 'sm_100', 'sm_120']
  • Entrypoint pre-flight: Unsloth container: 8 GPU(s). Primary: NVIDIA B200 sm_100 bf16=True
  • Imports: unsloth → unsloth_zoo → xformers → bnb → transformers → trl → peft (no order violations)
  • Model load: unsloth/Llama-3.2-1B-Instruct-bnb-4bit loaded in 4-bit, 16 QKV + 16 O LoRA layers patched
  • Training: 5 LoRA steps completed, loss decreased monotonically

Unsloth's own banner inside the container:

NVIDIA B200. Num GPUs = 8. Max memory: 178.353 GB. Platform: Linux.
Torch: 2.10.0+cu128. CUDA: 10.0. CUDA Toolkit: 12.8. Triton: 3.6.0
Bfloat16 = TRUE. FA [Xformers = 0.0.34. FA2 = False]

Loss progression

Bit-for-bit identical to the internal validation on a GCP B200 (different host, same image):

step 0  loss=4.1108
step 1  loss=4.0617
step 2  loss=3.9786
step 3  loss=3.8695
step 4  loss=3.7511

Bugs caught and fixed during validation

Each was a real defect surfaced only by running the image on a fresh host:

  1. `docker build --progress` rejected by the legacy builder → require buildx (c6d9216, 56d2701)
  2. Pre-flight missed `permission denied` on docker socket → check `docker info` exit code, print `usermod -aG docker` recipe (23a5b43)
  3. PEP 668 rejected the system-wide `pip install -U pip uv` bootstrap → dropped, the venv self-bootstraps via ensurepip (fd55ed0)
  4. `smoke_test.py` imported `unsloth_zoo` before `unsloth` → reorder (00cbc82)
  5. Triton's lazily-built CudaUtils C extension needed gcc + python3-dev at runtime → add to runtime stage (1cdc5f1)

Full gpt-oss-20B fine-tuning notebook run still pending; will post follow-up.

…ia.com/cuda/gpus

TORCH_CUDA_ARCH_LIST now covers the full set of compute capabilities
NVIDIA publishes on https://developer.nvidia.com/cuda/gpus for x86_64
hardware, from Turing onward:

  sm_75    Turing       T4, RTX 20-series, Quadro RTX
  sm_80    Ampere DC    A100, A30
  sm_86    Ampere       A40, RTX A6000, RTX 30-series
  sm_89    Ada          L4, L40, L40S, RTX 40-series
  sm_90    Hopper       H100, H200, GH200
  sm_100   Blackwell DC B100, B200, GB200
  sm_103   Blackwell DC B300, GB300
  sm_120   Blackwell    RTX 50-series, RTX PRO 6000 Blackwell
  sm_121   Blackwell    GB10 (DGX Spark)

with +PTX on the highest entry so future arch revisions can JIT.

Setting TORCH_CUDA_ARCH_LIST only affects nvcc invocations for any
source build the user adds on top of this image (e.g. flash-attn, a
custom CUDA op). The prebuilt cu128 wheels already include SASS for
sm_70/75/80/86/90/100/120 (verified at build time via
torch._C._cuda_getArchFlags()). Ada (sm_89), B300 (sm_103) and DGX
Spark (sm_121) GPUs run via JIT-PTX from the nearest available arch.

Jetson archs (sm_87 Orin, sm_110 Thor) are intentionally NOT included
-- they require aarch64 wheels and this image is linux/amd64 only.

Also lower the entrypoint's compute-capability gate from sm_80 to
sm_75. Turing GPUs work, with the caveat that bfloat16 is unavailable;
the entrypoint prints a NOTE in that case so Unsloth's fp16 fallback
isn't a surprise.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: dde5170e7a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread docker/entrypoint.sh
("sm_120", "Blackwell", "RTX 50-series, RTX PRO 6000 Blackwell"),
("sm_121", "Blackwell", "GB10 (DGX Spark)"),
)
if major < 7 or (major == 7 and minor < 5):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Reject pre-Ampere GPUs in entrypoint check

The entrypoint claims to fail fast on GPUs older than Ampere, but this condition only rejects devices below sm_75, so Turing (sm_75) is allowed through. That contradicts both the script’s own header (Unsloth requires sm_80+) and smoke_test.py, which exits on any cap[0] < 8; on T4/RTX20 hosts the container starts successfully and then fails later during validation/workloads instead of giving the intended immediate, actionable error.

Useful? React with 👍 / 👎.

HF Hub does not act as a generic OCI registry for arbitrary Docker
images -- the registry.hf.space endpoint only serves images that
Spaces have built, not images pushed by `docker push`. So we cannot
do `docker push huggingface.co/user/repo:tag` for an Unsloth image.

For cross-host testing where we want one canonical place to pull
from (and Docker Hub credentials are not yet configured), wrap the
manual flow into push/pull-shaped commands:

  hf_push.sh: docker save | pigz | huggingface-cli upload
  hf_pull.sh: huggingface-cli download | gunzip | docker load

This is approximation, not real OCI semantics -- every push uploads
the full ~4 GB blob, no layer dedup, no manifest negotiation. Good
for testing across A100 / H100 / RTX 6000 boxes; the real release
should go through .github/workflows/docker-publish.yml to Docker Hub,
which gets layer dedup, multi-arch manifest support, and standard
`docker pull` UX for users.

Usage:
  bash docker/hf_push.sh unsloth-blackwell:test danielhanchen/unsloth-blackwell-docker
  bash docker/hf_pull.sh danielhanchen/unsloth-blackwell-docker unsloth-blackwell-test.tar.gz unsloth-blackwell:test

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4bfb4b891a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread .github/workflows/docker-publish.yml Outdated
with:
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
tags: |
type=raw,value=latest,enable={{is_default_branch}}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Prevent workflow_dispatch builds from overwriting latest

This tag rule publishes latest whenever the run is on the default branch, but workflow_dispatch also allows overriding UNSLOTH_REF/UNSLOTH_ZOO_REF to arbitrary refs. A manual test run from main can therefore push a non-release image as latest, which makes downstream users pull an unintended build. Gate latest to trusted release flows (or only when both refs are the defaults) to avoid accidental retagging.

Useful? React with 👍 / 👎.

Comment thread docker/hf_pull.sh
set -euo pipefail

REPO="${1:?usage: hf_pull.sh <hf_repo> [<blob>] [<verify_tag>]}"
BLOB="${2:-unsloth-blackwell.tar.gz}"

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Match hf_pull default blob name to push naming

hf_push.sh uploads archives as <image-name>-<tag>.tar.gz, but hf_pull.sh defaults to unsloth-blackwell.tar.gz. If users run the documented short form bash docker/hf_pull.sh <hf_repo> after a normal push, the download will target a filename that was never uploaded and fail. Keeping both scripts on the same default naming convention avoids this broken default path.

Useful? React with 👍 / 👎.

…face-cli`

In huggingface_hub >= 0.27 the `huggingface-cli` binary is deprecated
and prints a "Use hf instead" notice then exits without doing the
operation. The previous wrappers ran `huggingface-cli upload/download`
silently, treated the deprecation exit as success, and uploaded
nothing.

Detect the new `hf` binary first and use that. If only the legacy
`huggingface-cli` is on PATH (older installs), fall back with a WARN
so users know the failure mode if anything goes sideways.

Also: hf_pull.sh now asserts the downloaded file is non-empty
(`test -s`) so we catch silent download failures before the
`docker load` step.
…t call

`jupyter nbconvert --to script nb.ipynb --output nb 2>/dev/null` was
silently exiting 0 without producing the output file in some
environments (likely because jupyter/jupyter_core wasn't on PATH or
nbconvert's --output handling differed across versions). The 2>/dev/null
hid the underlying error, and `set -e` did not catch the missing-output
case because nbconvert itself returned 0.

Switch to a direct nbformat-based conversion:

  pip install -q nbformat
  python -c "import nbformat; nb=nbformat.read('nb.ipynb', as_version=4);
             code='\n\n'.join(c.source for c in nb.cells if c.cell_type == 'code')
             open('nb.py','w').write(code + '\n')"

Smaller dep set, no shell-out to a jupyter wrapper script, and an
explicit `test -s nb.py` afterwards catches any silent failure
before downstream steps try to read the file.

Reproduces the failure on RTX PRO 6000 Blackwell (sm_120, docker
29.2.1, ubuntu 24.04) where nbconvert's CLI silently no-op'd.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8344fa0a56

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread docker/Dockerfile Outdated
TRITON_CACHE_DIR=/workspace/.cache/triton \
# Keep the arch list visible at runtime in case the user source-builds anything
# extra inside the container (e.g. a custom CUDA op).
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0;10.0;12.0+PTX"

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Include sm_75 in runtime arch list

The runtime image advertises support from Turing onward (see entrypoint.sh supported list), but TORCH_CUDA_ARCH_LIST here starts at 8.0. Any CUDA extension compiled inside the running container (the exact use case this env var comment describes) will be built without sm_75, so on T4/RTX20 hosts those kernels can fail at runtime with “no kernel image is available” despite the container claiming that architecture is supported.

Useful? React with 👍 / 👎.

The previous nbformat-based conversion dumped raw cell.source for every
code cell. The gpt-oss-20B notebook's first cell uses Jupyter !shell
magic to install dependencies:

  !pip install --upgrade -qqq uv
  !uv pip install -qqq ... \
  git+https://github.com/triton-lang/triton.git@0add68... ...

Dumped verbatim, the `@0add68...` token tripped the Python parser with
"SyntaxError: invalid decimal literal" before training could even start.

The container already has unsloth, triton, transformers, etc. baked in,
so we don't need the notebook's install cell. Skip any cell whose source
contains pip/install markers, and comment out stray !cmd / %magic lines
in any other cells. Then assert nb.py parses with ast.parse() before
trying to run it -- catches conversion failures up front instead of at
training time.

Reproduces on RTX PRO 6000 Blackwell (sm_120, fresh Docker 29.2.1
host) where the previous conversion produced an invalid nb.py.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 391532c031

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread docker/test_locally.sh Outdated
Comment on lines +122 to +123
git clone --depth 1 -b docker-blackwell-build \
https://github.com/unslothai/unsloth.git /tmp/unsloth-pr 2>&1 | tail -3

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Clone stable ref in fallback build-context path

When this script is run outside the repo tree, the fallback path clones a hardcoded docker-blackwell-build branch. That branch name is PR-specific and can disappear after merge, so the fallback clone will fail and Block 2 cannot build at all. This breaks the script’s advertised “clone if needed” flow for users validating from a clean host; use a stable default ref (or a configurable ref input) instead of a transient PR branch.

Useful? React with 👍 / 👎.

@danielhanchen

Copy link
Copy Markdown
Member Author

Cross-host validation #2: sm_120 (RTX PRO 6000 Blackwell, 96GB) — full gpt-oss-20B fine-tuning

End-to-end validated on a fresh GCP RTX PRO 6000 Blackwell Server Edition host. Image pulled from HF (huggingface.co/danielhanchen/unsloth-blackwell-docker), loaded into Docker, and exercised with both the smoke test AND the full gpt-oss-20B fine-tuning notebook.

Setup

  • Host: GCP, NVIDIA RTX PRO 6000 Blackwell Server Edition (sm_120), 97887 MiB, driver 580.126.09
  • Docker: 29.2.1, buildx 0.30.1
  • Image: identical `unsloth-blackwell:test` built on the AWS B200, distributed via HF Hub

Smoke test (5-step LoRA on Llama-3.2-1B-bnb-4bit)

```
Unsloth container: 1 GPU(s). Primary: NVIDIA RTX PRO 6000 Blackwell Server Edition sm_120 bf16=True
arches: ['sm_70', 'sm_75', 'sm_80', 'sm_86', 'sm_90', 'sm_100', 'sm_120']
step 0 loss=4.0946
step 1 loss=4.0556
step 2 loss=3.9781
step 3 loss=3.8706
step 4 loss=3.7453
=== all checks passed ===
```

Loss is ~0.4% off the B200 (sm_100) run (4.1108 → 3.7511); expected because sm_100 vs sm_120 Triton kernels produce slightly different bf16 rounding paths, deterministically.

Full gpt-oss-20B fine-tuning (10 LoRA steps, MXFP4 + MoE)

10 SFT steps on `HuggingFaceH4/Multilingual-Thinking` with the gpt-oss-20B MXFP4 model. Real workload, real MoE expert LoRA, real Harmony format inference at multiple reasoning_effort levels:

```
==((====))== Unsloth 2026.5.6: Fast Gpt_Oss patching. Transformers: 5.5.0.
\\ /| NVIDIA RTX PRO 6000 Blackwell Server Edition. Num GPUs = 1.
Max memory: 94.971 GB. Platform: Linux.
O^O/ \/ \ Torch: 2.10.0+cu128. CUDA: 12.0. CUDA Toolkit: 12.8. Triton: 3.6.0
\ / Bfloat16 = TRUE. FA [Xformers = 0.0.34. FA2 = False]
"-
___-"

Unsloth: Detected MoE model with num_experts = 32 and target_modules = [...].
Enabling LoRA on MoE parameters: ['mlp.experts.gate_up_proj', 'mlp.experts.down_proj']

Trainable parameters = 3,981,312 of 20,918,738,496 (0.02% trained)
Num examples = 933 | Num Epochs = 1 | Total steps = 10
Batch size per device = 1 | Gradient accumulation steps = 4

step 1: loss=1.071 grad_norm=2.805
step 2: loss=1.633 grad_norm=3.155
step 3: loss=1.053 grad_norm=2.815
step 4: loss=0.835 grad_norm=2.14
step 5: loss=1.363 grad_norm=2.239
step 6: loss=0.936 grad_norm=1.645
step 7: loss=0.958 grad_norm=1.641
step 8: loss=1.203 grad_norm=1.843
step 9: loss=1.296 grad_norm=2.013
step 10: loss=0.952 grad_norm=1.679

train_runtime = 141.7s (2.36 minutes)
peak reserved memory = 12.625 GB / 94.971 GB (13.3%)
```

Post-train inference at `reasoning_effort="medium"` and `"high"` produced coherent French reasoning output via the Harmony format -- confirming MXFP4 weights are loading correctly, MoE expert routing works, and the Triton kernels JIT-compile for sm_120 at first use.

What this validates

  • cu128 wheel's sm_120 SASS works on real sm_120 hardware
  • xformers + bnb + triton fat binaries run cross-arch (built for the same image that worked on B200 sm_100)
  • MXFP4 quantization (`unsloth/gpt-oss-20b`, no bnb-4bit variant) works inside the container
  • MoE expert LoRA targeting (`mlp.experts.gate_up_proj`, `mlp.experts.down_proj`) works
  • The whole `docker save → HF Hub → docker pull → docker load → docker run` flow works for a 13 GB image
  • Triton 3.6.0 + sm_120 first-run JIT cost is acceptable (the 1st training step is slow at 105s as Triton compiles; subsequent steps drop to ~3-7s/step)

Additional bugs caught and fixed during this validation

  • `huggingface-cli` is deprecated in huggingface_hub >= 0.27, silently exits without doing the upload/download. Switched `hf_{push,pull}.sh` to the new `hf` CLI (7354642).
  • `jupyter nbconvert` was silently failing in the container. Replaced with direct `nbformat`-based conversion (8344fa0).
  • `nbformat` conversion dumped raw `!pip install` shell magic from notebook install cells, breaking Python parse. Added install-cell skip + `ast.parse()` assertion (391532c).

Make the docker image multi-arch so DGX Spark (GB10, sm_121, aarch64) and
the Grace-Hopper / Grace-Blackwell SoCs (GH200 arm64, GB200 arm64) pull a
natively-built arm64 child from the same manifest. Runtime emulation is
NOT involved -- QEMU is used only for the cross-compile step on x86_64
CI runners; consumers on aarch64 hosts get a normal arm64 image and CUDA
works as on any other host.

Dockerfile:
  * ARG TARGETARCH; switch unsloth extras between cu128-ampere-torch2100
    (amd64, with xformers) and huggingface (arm64, no xformers -- there
    is no cu128 aarch64 xformers wheel as of 0.0.34, so we fall back to
    Unsloth's native SDPA path; ~5-10% slowdown but functionally complete).
  * Build-time torch._C._cuda_getArchFlags() assertion: amd64 still
    requires sm_120, arm64 accepts sm_120 or sm_121.
  * Same TORCH_CUDA_ARCH_LIST on both arches; nvcc emits whatever's listed.

docker/setup_qemu.sh (new):
  One-time host setup -- registers binfmt_misc handlers via
  tonistiigi/binfmt and creates a 'unsloth-multiarch' docker-container
  buildx builder. Required only on x86_64 build hosts targeting arm64.

docker/test_locally.sh:
  --platform amd64|arm64 flag. Cross-builds verify QEMU is registered,
  then build through the in-image arch-flags assertion. Smoke + notebook
  blocks auto-skip when image arch != host arch (CUDA cannot run under
  user-space QEMU + nvidia-container-toolkit cannot bridge a QEMU guest
  to a real GPU).

.github/workflows/docker-publish.yml:
  platforms: linux/amd64,linux/arm64 (single manifest, two children).
  Timeout bumped 60 -> 150 min for the slower arm64-under-QEMU leg.
  docker/setup-qemu-action@v3 with platforms: arm64 (was implicit before).
@danielhanchen

Copy link
Copy Markdown
Member Author

DGX Spark / linux/arm64 support added via QEMU at build time (e7cfcea).

The image is now multi-arch: one Docker manifest with linux/amd64 + linux/arm64 children. docker pull unsloth/unsloth:latest on x86_64 hosts gets the amd64 layer; on DGX Spark / Grace / Grace-Hopper it gets the arm64 layer natively, with normal CUDA access. QEMU is only used at build time on the x86_64 CI runner -- never at runtime, where it would break CUDA.

What changed:

  • docker/Dockerfile: ARG TARGETARCH switches the unsloth extras between cu128-ampere-torch2100 (amd64, with xformers) and huggingface (arm64, falls back to Unsloth's native SDPA path -- xformers does not yet publish a cu128 aarch64 wheel as of 0.0.34). Build-time torch._C._cuda_getArchFlags() assertion now accepts sm_120 or sm_121 on arm64.
  • docker/setup_qemu.sh: one-time host setup that registers binfmt_misc handlers via tonistiigi/binfmt and creates a multi-arch buildx builder. Only needed when building arm64 on an x86_64 host.
  • docker/test_locally.sh: --platform amd64|arm64 flag. When the target arch differs from the host arch, the smoke test and notebook blocks auto-skip with a clear warning (CUDA does not work under QEMU runtime emulation and nvidia-container-toolkit cannot bridge a QEMU guest to a real GPU). The build itself still runs through the arch-flags assertion, so PRs get build-time validation of the arm64 leg even on amd64 runners.
  • .github/workflows/docker-publish.yml: platforms: linux/amd64,linux/arm64 (single manifest). Timeout bumped 60 -> 150 min for the slower arm64-under-QEMU leg.

Arm64 GPU SoCs supported by the new variant:

Compute Cap SoC Examples
sm_90 Grace-Hopper GH200
sm_100 Grace-Blackwell GB200
sm_121 Blackwell + Grace GB10 (DGX Spark)

End-to-end arm64 validation on actual hardware is pending -- the build-time assertion exercises the wheel resolution and the cu128 aarch64 fat binary, but the final proof is docker run --gpus all on a real DGX Spark / GB200. That can run once we have access; the published manifest will already include the arm64 child for early users to test.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e7cfceadab

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread docker/smoke_test.py Outdated
import unsloth_zoo

print(f"unsloth_zoo {unsloth_zoo.__version__}")
import xformers

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Gate xformers import by platform in smoke test

The arm64 image path intentionally omits xformers (UNSLOTH_EXTRA="huggingface" in docker/Dockerfile) because no cu128 aarch64 wheel is expected, but check_imports() unconditionally does import xformers here. On arm64 runs this makes /workspace/smoke_test.py fail before training, so the published arm64 variant cannot pass the repository’s own runtime validation despite being a supported target.

Useful? React with 👍 / 👎.

Comment thread docker/smoke_test.py
print(f"cuda build {torch.version.cuda}")
print(f"arches {arches}")
assert "sm_100" in arches, f"sm_100 missing: {arches}"
assert "sm_120" in arches, f"sm_120 missing: {arches}"

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Accept sm_121 in smoke-test arch validation

This assertion hard-requires sm_120, but the same commit’s Dockerfile build validation explicitly treats arm64 Blackwell as valid when either sm_120 or sm_121 is present. As written, a valid arm64 build that reports only sm_121 will fail the smoke test with a false negative, even though the image is intended to support GB10/DGX Spark.

Useful? React with 👍 / 👎.

Sync the docker image branch with main (15 commits) so the branch's Python
files carry main's current formatting and pre-commit.ci runs cleanly on a
non-stale checkout.
Post-merge ruff-format-with-kwargs pass (the pre-commit.ci hook) on the merged
loader.py; whitespace only, no logic change.
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Jul 6, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Jul 6, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Jul 6, 2026
…6681)

* docker: Colab-grade JupyterLab and Studio UX for the Blackwell image

Stacks a Colab-like JupyterLab and Studio experience on top of the
existing Blackwell image. Additive only: the training stack, CUDA/torch
pinning, and the Studio/JupyterLab/sshd service trio are unchanged.

JupyterLab labextension (prebuilt in a throwaway builder stage, so the
runtime image stays Node-free):
  - Unsloth Dark (Monokai) theme, adaptive light/dark by system preference
  - Colab-style ArrowDown/Up cell navigation
  - top-bar Unsloth logo (stock Jupyter logo disabled and locked)
  - #@title lines render as collapsible Heading-2 form bars
  - Ctrl+A in a cell output selects only that output, not the whole
    notebook (the old behaviour ran notebook:select-all and was laggy)
  - right activity bar hidden by default
  - overrides.json: per-cell run button without auto-advance, labeled
    Restart and Run All, windowing off so collapsing an output does not
    snap to the cell top, news/update prompts suppressed

Studio and login branding: Unsloth favicon, page logo, and a dark
Unsloth login page that rotates through the curated Studio sloth
stickers (fail-soft to the logo).

Notebook organization and Colab compatibility (base image):
  - categorized folder view built from relative symlinks mirroring the
    README sections, rebuilt each boot; real .ipynb files never moved,
    and the symlink tree is invisible to the sync state machine
  - AMD-* notebooks shown only on an AMD/HIP host (autodetected)
  - Docker-only strip of the Colab "Run all on Colab" intro sentence
    from unedited notebooks (upstream notebooks unchanged)
  - hoist %%capture above a leading #@title form so the cell runs
  - the per-cell transformers-sidecar log is silent unless
    UNSLOTH_ENABLE_LOGGING=1

Dependency pinning and naming: the curated notebook extras are pinned to
their resolved versions for reproducible rebuilds; decord is split into
its own fail-soft install (no aarch64 wheel). The lean base image is
renamed from :base to :core.

Adds tests/validate_studio_features.py, a static self-test for the
labextension plugins, overrides keys, and branding wiring.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* docker: address review feedback on the JupyterLab/Studio UX

- unsloth_nb_view.py: rebuilding the categorized view no longer deletes
  user files. The view is also JupyterLab's landing dir, so a user may
  save real notebooks there; _clear_view now unlinks only the symlinks we
  own and removes only folders that end up empty, leaving regular files
  in place. It also tests islink before isdir, so a view that is itself a
  symlink to a directory is unlinked instead of being walked into (which
  would have wiped the symlink target).

- studio_launch.sh: derive the landing URL and preferred_dir from
  UNSLOTH_NOTEBOOKS_VIEW_DIR / UNSLOTH_SKIP_NOTEBOOK_VIEW, the same env
  the sync script uses, instead of hard-coding /workspace/Unsloth
  Notebooks. A relocated or disabled view no longer opens JupyterLab on a
  missing folder; it falls back to the default /lab over /workspace.

- Dockerfile.studio: the labext-builder stage now installs Node 20 from
  NodeSource. Ubuntu 24.04's distro nodejs is 18, below JupyterLab 4.6's
  declared Node >=20 engine. Node stays confined to the throwaway builder
  stage, so the runtime image is unchanged.

- .dockerignore: explicitly allowlist jupyter/install_sloth_stickers.py
  alongside its sibling jupyter assets, rather than relying on the
  directory re-inclusion.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* docker: publish lean image as :core and full image as :studio

Complete the base->core (and "studio as studio") tag rename so the publish
workflow matches the user-facing helpers and the Dockerfile.studio header.

- The lean training image now publishes as :core (core-<tag>, core-nightly,
  core-sha-*); run.sh / docker_confirm.* already told users to pull :core, but
  docker-publish.yml still tagged it :base, so that pull would have 404'd. The
  per-arch digest artifacts are renamed to match.
- The full Studio image keeps :latest and gains a stable :studio alias, matching
  the Dockerfile.studio header.

Both the merge and post-publish smoke-test metadata blocks are updated together.
Internal "base image" wording (the layer Studio builds FROM) is left as-is.

* docker: address second-round review feedback on the JupyterLab/Studio UX

- studio_launch.sh: also gate the categorized-view landing URL on
  UNSLOTH_SKIP_NOTEBOOK_SYNC (the entrypoint skips building the view entirely in
  that mode), not just UNSLOTH_SKIP_NOTEBOOK_VIEW, so a no-sync container does not
  land on a missing folder.

- Dockerfile.studio: scope the sticker-install "|| echo" fallback to only the
  sticker step via a { ...; } group. It was attached to the whole branding &&
  chain, so a failure in a REQUIRED step (JS resolve, favicon/logo/login copy)
  was swallowed and the build continued with broken branding.

- unsloth_nb_view.py: when creating the categorized symlinks, only replace our
  own stale symlinks; if a real user file already occupies that name, keep it and
  skip the link instead of os.remove-ing it.

- overrides.json: drop doNotDisturbMode (it silenced ALL JupyterLab toasts,
  including kernel-restart / connection-drop feedback). The news/update prompts
  are already off via fetchNews / checkForUpdates.

- Dockerfile: keep decord mandatory on amd64 (fail the build on a missing or
  incompatible wheel) and only fail-soft on arm64/other arches that have no wheel.

- cellNav.ts: do not hijack ArrowUp/Down when focus is in an interactive output
  widget / form control, or while a completion popup is open, so ipywidgets
  controls and autocomplete at cell boundaries keep working.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* docker: keep Studio branding RUN free of comments inside the line continuation

Move the sloth-sticker fail-soft explanation above the RUN so no comment line
sits between backslash-continued commands. BuildKit strips such comments, but
keeping the RUN body a plain && chain removes the ambiguity for non-BuildKit
builders and static linters. The { ...; } fail-soft scoping is unchanged.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* docker: AGPLv3 attribution + integrity guard for the Studio/JupyterLab image

Make it obvious the image is built by Unsloth and hard to white-label out with a
shallow find-and-replace, and surface the AGPLv3 license + copyright in the UI.

Visible attribution (labextension):
- Help > "About Unsloth Docker Studio" dialog (about.ts): Unsloth logo, the
  AGPLv3 notice, "Copyright 2026-Present the Unsloth team", and source/website/
  license links. Added to the Help menu and the command palette.
- The JupyterLab loading splash is replaced with a spinning Unsloth logo
  (splash.ts, provides ISplashScreen; honors prefers-reduced-motion). The stock
  @jupyterlab/apputils-extension:splash is disabled+locked at build time, like
  the stock logo.
- AGPLv3 footer (license + copyright + links) on the branded login page.
- Labextension relicensed AGPL-3.0-only; SPDX headers on every source file.

Anti-tamper (no encoded/obfuscated strings -- plain readable text only; the one
data URI is the logo image):
- A canonical, plain-text attribution set lives in unsloth_branding.py with a
  TypeScript mirror (branding.ts) bundled verbatim into the labextension, so the
  phrase, copyright, links and plugin ids are spread across independent layers.
- unsloth_branding.py verifies all of these across the installed files (AGPLv3
  text, login footer, theme, labextension package + built bundle strings, logo,
  favicon) and fails loudly if any are missing. It runs at three layers:
  build time (fails the image build), the whole-container launcher
  (studio_launch.sh refuses to start), and as a jupyter_server extension
  (refuses to serve JupyterLab).
- tests/studio/test_branding_guard.py: positive + per-marker negative coverage,
  plus a check that no base64/decoder obfuscation crept into the attribution.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* docker: address #6681 review round 2 (colab magics, output select, branding guard)

- unsloth_colab_compat.py: only hoist a leading `%%` cell magic above the Colab
  `#@title` form for magics whose body runs as code (capture/time/bash/python/
  ...). Content magics (%%writefile, %%html, %%latex, ...) are left untouched so
  the form comment is never injected into the written file / rendered output.
- outputSelect.ts: stop trusting the text selection anchor to decide ownership
  of Ctrl/Cmd+A. A stale selection inside an output survives a click onto a
  command-mode cell or the file browser, which made select-all keep re-selecting
  the old output. Gate on the keystroke target or the last pointer-down (reset to
  null on any click outside an output) instead.
- unsloth_branding.py: also reject page_config.json that disables the Unsloth
  labextension or any of its plugin ids via disabledExtensions (dict or list
  form); that leaves the bundle on disk so the prior checks passed while the
  logo/About/splash attribution was stripped at load. Lock unsloth-jupyterlab in
  Dockerfile.studio as well (defense in depth), and add guard tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* labext: pin JupyterLab extension deps; confirm.ps1 /login probe

Pin the unsloth-jupyterlab npm deps to exact versions matching the baked
jupyterlab==4.6.0 (builder stays 4.5.9, its newest release) instead of floating
^/~ ranges, so the same commit always builds the same labextension bundle.
Also probe JupyterLab /login (not /api, which 403s behind a password hash) in
the Windows confirmation script.

* docker: categorize AMD/domain notebooks and wire the feature validation into CI

unsloth_nb_view.parse_readme only reset the folder section on level-3
(###) headings. The notebooks README carries level-1 domain headers
(# AMD Notebooks, # Kaggle Notebooks) with their own nb/*.ipynb link
tables and no intervening ###, so those notebooks were mis-filed under
the previous stale section (all 148 AMD notebooks landed in Other
Notebooks on an --amd build). Reset on any heading level and strip a
leading emoji/symbol run so the domain notebooks get their own clean
folder.

Also run tests/validate_studio_features.py explicitly in the repo CPU
job. It is named validate_* (not test_*) so pytest never collected it,
which meant a regression in the notebook view, Colab compat, strip,
JupyterLab defaults or login branding failed CI only when run by hand.

* labext: use caret ranges so jlpm dedups JupyterLab/Lumino singletons

The exact pins introduced earlier (@jupyterlab/* 4.6.0, @lumino/widgets
2.8.0, @jupyterlab/builder 4.5.9) break the Dockerfile.studio
labext-builder stage. Exact-pinning the framework packages defeats
jlpm's (yarn classic) hoisting: transitive @jupyterlab deps request
caret ranges that resolve to newer patch releases (e.g. @jupyterlab/
notebook pulls @jupyterlab/cells ^4.6.0 -> a newer patch), so jlpm
installs a second nested copy alongside the exact top-level one. Two
copies of @jupyterlab/cells and @lumino/widgets in the tree produce
TS2345 "not assignable" errors (protected-member/identity mismatch)
and the build fails.

Caret ranges let jlpm collapse every @jupyterlab and @Lumino package
to a single hoisted copy, which is required for a JupyterLab prebuilt
(federated) extension: at runtime those packages are shared singletons
provided by the host JupyterLab, so the build-time versions only need
to type-check against one consistent tree, not match an exact runtime
patch. This is the version set the published image was built and
validated with end to end.

Verified by building the labext in isolation against the base image
(Node 20 + bundled jlpm): caret ranges build clean (webpack compiled
successfully); the exact pins fail with the duplicate-package TS
errors.

* ci(studio-backend): trigger on docker/** so the JupyterLab feature validation guards docker-only changes

The 'Docker JupyterLab/notebook feature validation' step runs
tests/validate_studio_features.py, which checks docker/jupyter (the
labextension, overrides.json, login branding) and the docker notebook
helpers. The pull_request paths filter listed studio/unsloth/tests but
not docker/**, so a PR that only touches docker/ would skip that step
and a regression in those files could pass CI. Add docker/** so the
validation runs whenever the files it checks change.

* jupyter: center the login card and place the attribution below it

#site was a flex container using the default row direction with two children
(the login card and the AGPLv3 attribution), so they rendered side by side:
the card sat left of centre and the attribution floated up to the top-right.
Stack them in a column so the card is horizontally centred and the attribution
sits below it as a footer, matching the intended single-column layout.

* jupyter: refresh Studio attribution, About dialog and loading splash

- Attribution now reads 'Built by the Unsloth team' with a single Apache 2.0 /
  AGPLv3 license link (to the repo license section) on the login page and in the
  About dialog, replacing the plain 'Built by Unsloth. Licensed under the GNU
  AGPLv3.' line. The integrity guard, its canonical PHRASE and the branding tests
  are updated to match.
- About dialog: left-align the link rows so the labels line up instead of each
  row centering independently; add an 'Unsloth Reference' link to the docs, and a
  Licenses section listing Unsloth Studio (AGPLv3) and Unsloth Core (Apache 2.0)
  alongside the full license link.
- Loading splash now reads 'Loading Unsloth Docker' instead of the attribution
  label, via a dedicated SPLASH_LABEL constant.

* docker: document the branding attribution as an AGPLv3 Section 7 notice

Add docker/NOTICE and docker/jupyter/BRANDING.md so the Unsloth attribution that
unsloth_branding.py enforces is also a written license condition, not only a
build check. docker/NOTICE designates the attribution (the "Built by the Unsloth
team" label, the copyright line, the license notice, the logo and theme, and the
Help > About links) as required Appropriate Legal Notices under AGPLv3 Section
7(b), referencing /studio/LICENSE.AGPL-3.0 and /LICENSE. BRANDING.md is a
human-readable note next to the guard describing what must stay, where it lives
and how it is enforced.

* ci(studio-backend): restore docker/** trigger path

The docker/** pull_request path added in b558bc7 was dropped by a later
rebase, so the "Docker JupyterLab/notebook feature validation" step (which runs
tests/validate_studio_features.py against docker/jupyter branding and notebook
helpers) no longer ran on PRs that only touch docker/. Re-add docker/** so a
docker-only change is validated on the PR rather than only after merge to main.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Jul 6, 2026
… to amd64

Base image (torch 2.11.0):
  - amd64 unsloth extra: cu128-ampere-torch2100 -> cu128-ampere-torch2110.
    The old extra pulls xformers 0.0.34, which hard-pins torch==2.10.0 and
    conflicts with the torch==2.11.0 held throughout the build; the torch2110
    family pulls xformers 0.0.35 (no torch pin) and resolves cleanly. This
    needs an unsloth carrying the torch2110 CUDA extras on main, so merge the
    torch2110 extras PR first (default UNSLOTH_REF=main).
  - notebook-deps assertion: startswith('2.10.0') -> '2.11.0' so the layer
    actually verifies the torch it now installs.
  - refresh the torch2100/xformers 0.0.34 references in the surrounding
    comments to the torch2110/0.0.35 line.

sm_103 (B300/GB300) JIT override (Codex item):
  The cu13 NVRTC/ptxas override was arm64-only (sm_121), and its comment
  claimed triton 3.6.0 bundles cu13 ptxas and set TRITON_PTXAS_PATH -- neither
  was true: triton 3.6.0's bundled ptxas is CUDA 12.8 (V12.8.93, tops out at
  sm_120) and TRITON_PTXAS_PATH was never set. So sm_103 (amd64) and even
  sm_121 (arm64) Triton JIT were unfixed.

  Run the cu13 install on both arches and actually wire the ptxas override:
    - NVRTC swap (cu13 libnvrtc.so.13 over torch's bundled cu12.8 .so.12) now
      runs on amd64 too.
    - ENV TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas routes every Triton
      JIT through the cu13 ptxas. Global rather than per-arch is safe: cu13.0
      ptxas spans sm_70..sm_121 (verified: Volta/Turing/Ampere/Hopper through
      Blackwell), so no regression for the older GPUs in the arch list.

  Verified on amd64 in the built base image: cuda-nvrtc-13-0/cuda-nvcc-13-0
  install cleanly from the base's CUDA repo, ptxas lands at
  /usr/local/cuda-13.0/bin/ptxas (V13.0.88) and libnvrtc.so.13 at
  /usr/local/cuda-13.0/lib64/. The sm_103/sm_121 runtime path itself is not
  hardware-tested (no such GPU on hand); precompiled SASS still covers both
  via sm_100/sm_120 forward-compat, so only JIT-heavy paths rely on this.
…before CUDA dedup

The Studio build symlinks the Studio venv's CUDA libs onto the base venv's
copies to reclaim ~3.7GB. That is only safe when both venvs run the same torch,
but the pre-dedup guard only checked the CUDA family (endswith('+cu128')). A
Studio venv that installed torch 2.10.0+cu128 (an installer capped below the
base's 2.11.0, or a build-time nvidia-smi fallback) would pass that check yet
mismatch the base's 2.11.0+cu128, and the dedup would link incompatible libs.

Capture the base venv's torch from its metadata and assert the Studio venv torch
equals it exactly (version and family) before the dedup runs, so a mismatch
fails the build loudly instead of silently linking skewed CUDA libs. Comparing
to the base venv also avoids hardcoding the version here. The Studio venv reaches
torch 2.11.0+cu128 via the installer's UNSLOTH_TORCH_INDEX_FAMILY=cu128 handling
and its CUDA torch spec allowing 2.11.x.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 386d3a7c74

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread .github/workflows/docker-publish.yml Outdated
CUDA_VERSION=12.8.1
UBUNTU_VERSION=24.04
PYTHON_VERSION=3.12
UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || github.sha || 'main' }}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Resolve UNSLOTH_REF before fan-out

When workflow_dispatch uses the default unsloth_ref=main or any branch name, this raw ref is fetched independently by each base matrix leg and again by the Studio build (UNSLOTH_STUDIO_REF uses the same expression). If that branch advances during the run, the amd64/arm64 children or the base and Studio layers can bake different Unsloth commits under one manifest/SHA tag; resolve the requested ref to one SHA in prepare and pass that output everywhere, like the zoo and notebooks refs already do.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d4dc8b6. Added a Resolve unsloth ref step to the prepare job that freezes the requested ref to one sha (same precedence as before: dispatch input, else pushed tag, else triggering sha, else main; ls-remote for a branch/tag), and rewired UNSLOTH_REF and UNSLOTH_STUDIO_REF to needs.prepare.outputs.unsloth_ref. Mirrors the existing zoo/notebooks resolution so all legs bake one commit.

Comment thread docker/Dockerfile.studio Outdated
Comment on lines +160 to +164
&& if [ "${TARGETARCH:-amd64}" = "arm64" ]; then \
for NVRTC_DIR in "${UNSLOTH_STUDIO_HOME}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do \
if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ]; then \
mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \
ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Swap the Studio NVRTC on amd64 too

On the full Studio image running on amd64 B300/GB300 (sm_103), this condition leaves the Studio venv's own nvidia/cuda_nvrtc/lib/libnvrtc.so.12 at CUDA 12.8. The base venv is swapped to cu13 and Triton gets a cu13 ptxas, but Studio's separate torch wheel can still load its bundled cu12.8 NVRTC via the wheel library path, so Studio-side NVRTC/Jiterator paths can fail on compute_103; apply the same symlink for amd64 or otherwise dedup/swap the Studio NVRTC there as well.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d4dc8b6. Dropped the arm64-only guard so the Studio venv NVRTC cu13 swap runs on both arches. amd64 sm_103 needs it too, the CUDA dedup never touches cuda_nvrtc, and the base cu13 layer installs cuda-nvrtc-13-0 on both arches so /usr/local/cuda-13.0/lib64/libnvrtc.so.13 exists regardless of TARGETARCH.

Comment on lines +65 to +66
"--upgrade-package",
"-P",

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Filter uv upgrade-package values

When a notebook uses uv's upgrade selector for a protected package, e.g. uv pip install -P torch peft or --upgrade-package=transformers, this value is treated as an opaque option value and kept verbatim instead of passing through _KEEP / the transformers marker. uv documents this flag as allowing upgrades for a specific package, so installing a kept target can still refresh torch or transformers and clobber the baked stack; filter these values the same way as normal requirement specs.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d4dc8b6. -P/--upgrade-package values now go through _KEEP: a protected value drops the flag+value pair (space and =forms), so uv pip install -P torch peft becomes uv pip install peft. Verified by running the shim on the exact example.

Comment on lines +49 to +53
_VALUE_FLAGS = {
"-r",
"--requirement",
"-c",
"--constraint",

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Pair editable flags with their targets

Because -e/--editable are not treated as value-taking flags (checked python -m pip install --help: -e, --editable <path/url>), the scanner can leave a dangling editable flag after dropping a protected target. For example, pip install -e git+https://github.com/unslothai/unsloth.git#egg=unsloth peft drops the VCS URL but keeps -e, then executes pip install -e peft, failing instead of installing the kept package; handle editable as a paired target and drop both flag and value when the target is protected.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d4dc8b6. -e/--editable is now a value-taking flag paired with its target; a protected editable drops both flag and value, so pip install -e git+...unsloth...#egg=unsloth peft becomes pip install peft with no dangling -e. Non-protected editables (e.g. -e ./mypkg) are unchanged.

Comment thread docker/unsloth_pip_shim.py Outdated
Comment on lines +102 to +112
if re.match(r"^[a-z]+\+", token) or "://" in token or token.startswith((".", "/")):
# A VCS / URL install can still name a protected package via the legacy
# `#egg=NAME` (or `&egg=NAME`) fragment, e.g.
# `git+https://github.com/unslothai/unsloth.git#egg=unsloth`. Pull that
# name out so _KEEP can drop it; otherwise the shim would exec the URL
# and reinstall a baked package into the venv. A non-protected egg name
# is returned too, but the caller keeps it as a normal target either way.
_egg = re.search(r"[#&]egg=([A-Za-z0-9][A-Za-z0-9._-]*)", token)
if _egg:
return _egg.group(1).lower().replace("_", "-") or None
return None # vcs / url / local path -> let it pass through

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Parse protected wheel filenames before URL passthrough

This URL/path passthrough still lets bare wheel installs for protected packages bypass _KEEP: a valid cell such as pip install https://download.pytorch.org/.../torch-2.11.0+cu128-...whl returns None here, is kept as a positional target, and later runs the real pip, reinstalling torch into the baked venv. This is separate from the fixed PEP 508 and #egg= cases; parse wheel basenames (and requirement-file wheel lines) for protected distribution names before allowing URL/local-wheel passthrough.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d4dc8b6. _canon now parses the PEP 427 distribution name from a wheel URL/path, so pip install https://.../torch-2.11.0+cu128-...whl drops instead of reinstalling the baked torch. Non-protected wheels still pass through verbatim.

danielhanchen and others added 2 commits July 8, 2026 05:02
…p-shim edges)

docker-publish.yml: freeze the requested unsloth ref to one sha in the prepare
job before the matrix fans out. UNSLOTH_REF / UNSLOTH_STUDIO_REF were raw
expressions re-evaluated per base arch leg and in the Studio build, so a mutable
branch (the workflow_dispatch default unsloth_ref=main) advancing during the run
could bake different unsloth commits under one manifest. Resolve once (same
precedence: dispatch input, else pushed tag, else triggering sha, else main;
ls-remote a branch/tag to a sha, mirroring the zoo/notebooks steps) and read
needs.prepare.outputs.unsloth_ref everywhere.

Dockerfile.studio: run the Studio venv NVRTC cu13 swap on both arches, not arm64
only. amd64 sm_103 (B300/GB300) needs cu13 NVRTC just as arm64 sm_121 does, and
the CUDA dedup never touches cuda_nvrtc, so an amd64 Studio venv would otherwise
keep its bundled cu12.8 libnvrtc and fail NVRTC/jiterator JIT on compute_103. The
base cu13 layer installs cuda-nvrtc-13-0 on both arches, so the target .so.13
exists here regardless of TARGETARCH.

unsloth_pip_shim.py: close three ways a protected package slipped past _KEEP.
Treat -e/--editable as a value-taking flag paired with its target and drop both
when the target is protected (was leaving a dangling -e that failed the cell);
filter -P/--upgrade-package values through _KEEP (a named baked package could be
refreshed while installing another target); and parse the PEP 427 distribution
name out of a wheel URL/path so a bare `pip install https://.../torch-...whl`
drops instead of reinstalling the baked torch. Non-protected editables, upgrade
selectors, and wheels are unchanged. Adds tests/python/test_unsloth_pip_shim.py
(18 regression tests, exec captured via a patched os.execv).
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Jul 8, 2026
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Jul 8, 2026

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d6e559008f

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread docker/unsloth_pip_shim.py Outdated
Comment on lines +253 to +254
new_line, rewrote, inc_rec, inc_drp = _rewrite_include(line, stripped, src_dir, _depth)
out.append(new_line)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Filter editable entries in requirement files

Fresh evidence beyond the fixed command-line editable case: when a notebook runs pip install -r reqs.txt and reqs.txt contains a valid editable requirement like -e git+https://github.com/unslothai/unsloth.git#egg=unsloth, this branch treats every - line as an option/include and preserves non-include lines unchanged. python -m pip install --help documents -e, --editable <path/url>, and pip accepts that form from requirement files, so protected editable requirements can still bypass _KEEP and reinstall the baked stack; classify -e/--editable requirement-file lines before appending them.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 251e3ed. An -e/--editable line inside a -r requirements file is now classified through _KEEP (new _parse_editable) and dropped when the target is protected, so -e git+...unsloth...#egg=unsloth in reqs.txt no longer reinstalls the baked stack. Non-protected editables are forwarded verbatim.

# as an opaque option and a `-r` file is never filtered -- and worse, it
# never counts as a target, so a cell whose only target is that file
# silently no-ops and installs nothing.
if tok.startswith("--") and "=" in tok:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Handle attached short requirement options

For valid pip forms such as pip install -rrequirements.txt and pip install -cconstraints.txt peft (the local pip parser accepts -rfile/-cfile, and python -m pip install --help shows these options consume a <file>), this scanner only handles the separated form or long --flag=value. The attached short token starts with - and falls through as an opaque option, so an -r...-only cell no-ops with has_target false, while -c... stays unfiltered and can still apply protected pins during resolution.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 251e3ed. Attached short forms -rreqs.txt / -cconstraints.txt / -epath / -Pname are split into flag+value and routed through the same requirement-file / constraint / _KEEP handling as the separated form, so an attached -r no longer no-ops and -c/-e/-P values are filtered.

Comment on lines +257 to +258
if inc_rec and not recorded:
recorded = inc_rec

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Do not record nested constraint pins as requests

Fresh evidence beyond the earlier constraint-file filtering fix: a requirements file that includes -c constraints.txt still propagates inc_rec here even when the nested include is a constraint, while the top-level -c path intentionally ignores _c_rec because a constraint is not a transformers install request. With reqs.txt containing -c constraints.txt plus a kept package, and the constraint pinning transformers==4.55.0, the shim writes the sidecar marker and the next model cell can activate that older transformers version even though it was only a resolver constraint; distinguish nested -c from nested -r before recording.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 251e3ed. A nested -c constraint include inside a -r file no longer propagates its transformers pin as a recorded request (a constraint is not an install request), mirroring the top-level -c path; the dropped-specs info still propagates.

Comment thread docker/Dockerfile.studio
&& git fetch -q --depth 1 origin "${UNSLOTH_STUDIO_REF}" \
&& git checkout -q FETCH_HEAD \
&& UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME}" \
UNSLOTH_TORCH_INDEX_FAMILY="${TORCH_FAMILY}" \

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Let the Studio install resolve torch 2.11 on cu128

Fresh evidence beyond the earlier Studio torch-thread reply: this build now passes UNSLOTH_TORCH_INDEX_FAMILY=cu128, but the install.sh --local it invokes still keeps CUDA on TORCH_CONSTRAINT="torch>=2.4,<2.11.0" and only relaxes that constraint for */rocm7.2, not cu128. The Studio venv therefore resolves a 2.10.x+cu128 torch, and the equality check a few lines later compares it with the base venv's 2.11.0+cu128 and aborts every Studio image build; thread an exact/base torch spec or relax the CUDA cu128 constraint before running the installer.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed, and this is provided by #6959 (branch studio-cuda-torch-2110): it adds a cu[0-9]*) case to install.sh widening cu128 to <2.12.0 (matched on _torch_index_leaf) plus the _CUDA_TORCH_PKG_SPEC repair-spec widening, so the Studio venv resolves 2.11.0+cu128 and the exact-match assert passes. The two must land together (#6959 with/before this PR).

Comment thread docker/Dockerfile Outdated
# above (triton 3.6.0's own ptxas is cu12.8, no sm_103/sm_121). Set globally,
# not per-arch: cu13.0 ptxas spans sm_70..sm_121 so it is correct for every GPU
# this image supports, and ENV cannot be made conditional per arch.
ENV TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Gate CUDA 13 JIT tools by driver/device

On otherwise-supported Ampere/Ada/Hopper/Turing hosts with 570–579 drivers, this global override sends every Triton JIT through CUDA 13 ptxas even though the entrypoint still advertises a 570+ driver floor except for sm_103/sm_121. NVIDIA's CUDA compatibility table lists CUDA 13.x as requiring a 580+ driver, so ordinary training JITs can fail on drivers the image currently accepts; enable the CUDA 13 ptxas/NVRTC swap only for sm_103/sm_121 or raise the documented/enforced driver floor.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 251e3ed. Confirmed: a cu13 cubin needs a >=580 driver to LOAD even when targeting an older arch (CUDA has forward, not backward, cross-major driver compat), so a global cu13 override breaks 570-579 hosts on Turing..sm_120. Removed the global ENV; select_cuda_jit_tools() in entrypoint.sh reads the device compute_cap and activates cu13 only for sm_103/sm_121 (which ship >=580 drivers), else keeps the bundled cu12.8 ptxas and restores the cu12.8 NVRTC in both venvs. 7 device-gating unit tests added.

…gate cu13 JIT tools)

unsloth_pip_shim.py: close three more ways a protected package slipped past
_KEEP. An editable line (-e/--editable <target>) inside a -r requirements file
is a real install target, so a protected editable there is now classified and
dropped like the command-line case (new _parse_editable). pip/uv accept the
attached short forms -rreqs.txt / -cconstraints.txt / -epath / -Pname as one
token; these were falling through as opaque options (so an attached -r-only cell
no-op'd and an attached -c/-e/-P value bypassed _KEEP), so the 2-char flag is now
split from its value and routed through the separated-form handling. And a nested
-c constraint inside a -r file no longer records its transformers pin as an
install request (a constraint is not a request; mirrors the top-level -c path).

entrypoint.sh / Dockerfile: gate the CUDA 13 ptxas + NVRTC to sm_103 / sm_121 at
runtime instead of a global build-time default. A cu13 cubin needs a >= 580
driver to LOAD even when it targets an older arch (CUDA has forward, not
backward, cross-major driver compatibility), but the image supports Turing..
sm_120 on a 570+ driver, so the previous global TRITON_PTXAS_PATH ENV + cu13
NVRTC symlink would break ordinary Triton/NVRTC JIT on 570-579 driver hosts. The
build still bakes cu13 (saving the cu12.8 NVRTC as .cu128.orig); a new
select_cuda_jit_tools() in the entrypoint reads the device compute_cap and only
activates cu13 for sm_103/sm_121 (which ship >= 580 drivers), otherwise leaving
Triton on its bundled cu12.8 ptxas and restoring the cu12.8 NVRTC in both the
base and Studio venvs. The base ENTRYPOINT runs for the Studio image too.

Adds 9 pip-shim regression tests and tests/sh/test_select_cuda_jit_tools.sh
(7 device-gating cases); registers the latter in CI and tests/run_all.sh.
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3bb40e47fe

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

)
if _dref:
return _dref.group(1).lower().replace("_", "-") or None
if re.match(r"^[a-z]+\+", token) or "://" in token or token.startswith((".", "/")):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Parse bare wheel filenames before passthrough

The fresh bypass is a bare relative wheel filename: pip accepts pip install torch-...whl from the current directory, but this guard only enters the wheel-basename parser for URLs/VCS or paths starting with .//. That token falls through as the normalized full filename, misses _KEEP, and the real pip can reinstall torch/unsloth/etc. into the baked venv; run the .whl parser before this guard or include bare *.whl tokens.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in b3649d4. _canon now parses a bare wheel filename (no ./ or / prefix and no scheme) via its PEP 427 basename before the plain-spec fallthrough, so pip install torch-2.11.0-...whl from the CWD resolves to torch and is dropped by _KEEP. Added test_bare_torch_wheel_filename_dropped / _in_subdir_dropped / _unprotected_kept.

Comment thread docker/entrypoint.sh Outdated
local cc="" nvrtc_dir orig
if command -v nvidia-smi >/dev/null 2>&1; then
cc="$( { nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true; } \
| head -n1 | tr -d '[:space:]' )"

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Check every visible GPU before restoring cu128 JIT tools

On mixed GPU hosts where nvidia-smi lists an H100/B200 before a B300/GB10, this line discards all but the first compute capability. Startup then takes the default branch, restores cu12.8 NVRTC and leaves Triton's ptxas unset even though a visible sm_103/sm_121 device still needs the CUDA 13 tools to JIT; workloads pinned to that later GPU hit the compute_103/121 failures this image is trying to avoid. Consider scanning all visible caps and enabling cu13 if any visible GPU requires it.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in b3649d4. select_cuda_jit_tools() now scans every visible compute_cap (a while-read over the full nvidia-smi list) and enables cu13 if ANY GPU is sm_103/sm_121, so a B300/GB10 sitting behind an H100/B200 still gets the cu13 tools. Those parts only ship on >=580 drivers, so the whole host tolerates cu13 cubins for every arch present. Added mixed-host tests (H100+B300, B200+GB10, and the order-reversed B300+H100).

Comment on lines +280 to +281
except OSError:
return path, None, [] # remote URL / unreadable -> let the real tool handle it

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Reject remote requirement files before pip sees them

For a notebook cell like pip install -r https://.../reqs.txt, pip accepts the URL form (the local pip install --help says the requirement target can be a file or URL), but this path treats the URL as unreadable and passes it through unchanged. Protected specs inside that remote file never go through _KEEP or the transformers sidecar marker, so the real pip can still reinstall torch/transformers into the baked venv; either fetch/filter URL requirement files or reject them in notebook mode.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in b3649d4. Remote (URL) -r/-c requirement/constraint files are now refused in shim mode: the top-level form in all three arg spellings (-r URL, --requirement=URL, -rURL) and nested -r/-c URL includes inside a local file, since their pins cannot be inspected before the real tool would fetch and install them. A remote -r as the sole target no-ops with no dangling flag; other targets on the same line still install. Added test_remote_requirement_url_* and test_nested_remote_include_dropped.

Comment thread docker/unsloth_pip_shim.py Outdated
dist = _whl.group(1).split("-", 1)[0].strip().lower().replace("_", "-")
if dist:
return dist
return None # vcs / url / local path -> let it pass through

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Drop protected VCS URLs without egg fragments

Fresh evidence beyond the fixed #egg= case: pip also accepts VCS project URLs without an egg fragment (local help lists VCS project URLs, and this repo recommends git+https://github.com/huggingface/transformers.git in unsloth/models/loader.py). Those tokens reach this fallback with _canon(...) == None, are kept as install targets, and can overwrite the baked transformers/unsloth stack without _KEEP or the marker ever firing; infer protected names from known repo basenames or block unclassified protected VCS URLs.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in b3649d4. _canon now infers the distribution from an egg-less VCS URL by its repo basename (git+https://github.com/huggingface/transformers.git -> transformers, unslothai/unsloth-zoo.git -> unsloth-zoo), stripping a trailing @ref and .git, so the egg-less form loader.py recommends is matched against _KEEP. A non-protected repo returns its basename and the caller keeps it as a normal target either way. Added test_vcs_url_without_egg_protected_dropped / _with_ref_dropped / _unprotected_kept.

continue
name = _canon(tok)
if name is None:
keep_args.append(tok) # bare flag, or a positional url / path / vcs

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Drop resolver-wide reinstall flags in notebook shim

In notebook cells such as pip install --force-reinstall peft or pip install --force-reinstall transformers==4.55.0 peft, this branch keeps the global reinstall flag after protected specs are dropped. Local pip install --help defines --force-reinstall as reinstalling all packages even if up-to-date, so the real pip can still replace protected dependencies pulled by the kept target (for example transformers/torch) even though direct protected requirements were filtered; strip these resolver-wide reinstall/ignore flags in shim mode or force dependency installation off for the kept command.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in b3649d4. Resolver-wide reinstall / ignore-installed switches (pip --force-reinstall, --ignore-installed, -I; uv --reinstall) are stripped in shim mode so they cannot rebuild already-satisfied baked deps pulled in by a kept target. The kept target still installs. Added test_force_reinstall_flag_stripped / _ignore_installed_short_flag_stripped / _uv_reinstall_flag_stripped.

Comment thread docker/unsloth_pip_shim.py Outdated
# package (e.g. `uv pip install -P torch peft`) lets an ordinary install target
# refresh that package and clobber the pinned stack. Filter its value through
# _KEEP too. Unlike -e it is not itself an install target (no has_target).
_UPGRADE_PKG_FLAGS = {"-P", "--upgrade-package"}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Filter uv reinstall-package selectors too

Fresh selector beyond the earlier -P fix: local uv pip install --help documents --reinstall-package <name> as reinstalling a specific package, but this protected-selector set only includes -P/--upgrade-package. As a result, uv pip install --reinstall-package=torch peft is passed through and can reinstall torch, while the separated form drops torch but leaves a dangling --reinstall-package that consumes the next kept target; route --reinstall-package values through the same _KEEP handling as upgrade-package.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in b3649d4. --reinstall-package is now in both _VALUE_FLAGS and _UPGRADE_PKG_FLAGS, so its value is filtered through _KEEP exactly like -P/--upgrade-package: a protected name drops the flag+value pair together (no dangling selector left to swallow the next target) in both --reinstall-package torch peft and --reinstall-package=torch peft. Added test_reinstall_package_protected_separated_dropped / _inline_dropped, _unprotected_kept and _transformers_pin_recorded.

@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Keep it up!

Reviewed commit: d34cb71189

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

danielhanchen and others added 2 commits July 8, 2026 07:16
Notebook pip/uv shim (docker/unsloth_pip_shim.py), all active only under
UNSLOTH_NB_SHIM=1:
  - Parse a bare wheel filename (torch-*.whl in the CWD, no ./ or / prefix) so
    it is matched against _KEEP instead of passing through as an opaque
    positional and reinstalling the baked torch.
  - Infer the distribution from an egg-less VCS URL by repo basename
    (git+https://github.com/huggingface/transformers.git -> transformers) so
    the egg-less form the repo itself recommends cannot clobber the baked stack.
  - Refuse remote (URL) -r/-c requirement/constraint files -- top-level and
    nested includes -- since their pins cannot be inspected before the real
    tool would fetch and install them.
  - Strip resolver-wide reinstall/ignore-installed switches
    (--force-reinstall, --ignore-installed, -I, uv --reinstall) so they cannot
    rebuild already-satisfied baked deps pulled in by a kept target.
  - Route uv --reinstall-package through the same _KEEP handling as
    -P/--upgrade-package (both attached and separated forms; no dangling flag).

Entrypoint (docker/entrypoint.sh): select_cuda_jit_tools() now scans every
visible GPU's compute_cap instead of only the first, so a datacenter Blackwell
(sm_103/sm_121) behind an H100/B200 still enables the cu13 JIT tools it needs.

Adds regression tests for each case (tests/python/test_unsloth_pip_shim.py,
tests/sh/test_select_cuda_jit_tools.sh).

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 167fdf26b9

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

_seg = _seg.strip().lower().replace("_", "-")
if _seg:
return _seg
return None # plain url / local path -> let it pass through

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Parse protected source archives before URL passthrough

When a notebook or requirements file names a protected package as a source archive instead of a wheel, for example pip install https://files.pythonhosted.org/.../unsloth-2026.7.1.tar.gz or ./torch-2.11.0.tar.gz, pip accepts <archive url/path> targets (checked local python -m pip install --help), but _canon() only extracts wheel basenames before returning None here. That token is kept and the real pip can reinstall unsloth/torch/transformers into the baked venv, bypassing _KEEP; parse sdist/zip basenames or reject protected archive URLs/paths before this passthrough.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 6a078b1. Added a _sdist_name() helper and wired it into both the URL/path branch and the bare-filename path of _canon, so a protected source archive (https://.../unsloth-2026.7.1.tar.gz, ./torch-2.11.0.tar.gz, .zip, .tar.bz2, etc.) is matched against _KEEP instead of passing through. It splits at the first hyphen before a digit so hyphenated names like flashinfer-python survive. Added test_sdist_url/bare/zip/hyphenated_protected_dropped and _unprotected_kept.

Comment thread docker/unsloth_pip_shim.py Outdated
Comment on lines +80 to +85
_REQ_FILE_FLAGS = {"-r", "--requirement"}
# Constraint files are not install targets, but pip applies their pins during
# resolution, so a `-c constraints.txt` that pins torch/transformers/etc. can
# still downgrade or reinstall a baked package when another target pulls it in.
# Filter protected packages out of them the same way as requirement files.
_CONSTRAINT_FILE_FLAGS = {"-c", "--constraint"}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Route uv plural requirements through the filter

When a notebook uses uv's documented plural flags, e.g. uv pip install --requirements reqs.txt or uv pip install --constraints constraints.txt peft (checked local uv pip install --help: -r, --requirements <REQUIREMENTS> and -c, --constraints <CONSTRAINTS>), these sets only recognize the singular pip-style names. The scanner therefore forwards those files to uv unfiltered, so protected pins inside can still reinstall or constrain torch/transformers despite the requirement/constraint filtering added for the other spellings.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 6a078b1. Added uv's plural long forms --requirements and --constraints to _REQ_FILE_FLAGS / _CONSTRAINT_FILE_FLAGS (and _VALUE_FLAGS), so uv pip install --requirements reqs.txt / --constraints constraints.txt go through the same protected-package filter as the singular pip-style names. Added test_uv_plural_requirements_filtered and test_uv_plural_constraints_filtered.

# guise of installing an unprotected package. The kept target still installs; its
# already-satisfied protected deps are left untouched. Per-package selectors
# (--reinstall-package / -P) are handled through _UPGRADE_PKG_FLAGS instead.
_REINSTALL_FLAGS = {"--force-reinstall", "--ignore-installed", "-I", "--reinstall"}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Strip eager upgrade flags before kept installs

For notebook cells like pip install -U --upgrade-strategy eager peft, these resolver-wide upgrade flags are still forwarded even though the local pip help says eager mode upgrades dependencies regardless of whether installed versions satisfy requirements. A kept target such as peft can then make pip upgrade protected dependencies like torch or transformers underneath the baked CUDA stack; strip -U/--upgrade with eager upgrade strategy (or otherwise force protected deps to stay installed) in the shim path too.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 6a078b1. --upgrade-strategy is now a value flag whose flag+value pair is dropped in shim mode (_DROP_VALUE_FLAGS), so pip install -U --upgrade-strategy eager peft falls back to pip's default only-if-needed and cannot eagerly rebuild already-satisfied baked deps like torch/transformers. The kept target (peft) still installs. Added test_upgrade_strategy_eager_dropped / _eager_inline_dropped / _only_if_needed_also_dropped.

Comment thread docker/entrypoint.sh Outdated
/opt/unsloth-venv/lib/python*/site-packages/nvidia/cuda_nvrtc/lib \
"${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do
orig="${nvrtc_dir}/libnvrtc.so.12.cu128.orig"
[[ -e "${orig}" ]] || continue

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Do not leave CUDA 13 NVRTC active for non-root users

When the container is started with --user on a non-sm_103/sm_121 host, this symlink restore cannot write into the root-owned venv and is silently ignored. Because the Dockerfile bakes libnvrtc.so.12 pointing at CUDA 13 first, supported Ampere/Ada/Hopper/Turing hosts on 570-579 drivers keep CUDA 13 NVRTC active even though the gate decided they need cu12.8, so later NVRTC/Jiterator paths can emit cubins the accepted driver cannot load; make cu12.8 the immutable default or run the switch before dropping privileges.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 6a078b1. Inverted the NVRTC default: the build now keeps cu12.8 as the immutable default (libnvrtc.so.12 -> .cu128.orig) and stages a .cu13 alias beside it; select_cuda_jit_tools retargets libnvrtc.so.12 -> .cu13 ONLY for sm_103/sm_121. Because the safe default needs no runtime write, a non-root docker run --user container -- which cannot rewrite the symlink -- now keeps cu12.8 (loadable on every 570+ driver) instead of being stuck on cu13. Applied to both the base and Studio Dockerfiles; entrypoint gating test updated to assert the cu12.8 default is left untouched for non-DC arches and switched to cu13 for sm_103/sm_121.

Notebook pip/uv shim (docker/unsloth_pip_shim.py, active only under
UNSLOTH_NB_SHIM=1):
  - Parse protected source archives (sdist/zip) by basename too, e.g.
    `pip install https://.../unsloth-2026.7.1.tar.gz` or `./torch-2.11.0.tar.gz`,
    mirroring the wheel-basename handling. A first-hyphen-before-digit split
    keeps hyphenated names like flashinfer-python intact.
  - Recognise uv's PLURAL long flags --requirements / --constraints, so those
    files go through the same protected-package filter as the singular names.
  - Drop --upgrade-strategy eager in shim mode so a kept target cannot eagerly
    rebuild already-satisfied baked deps (falls back to pip's only-if-needed).

NVRTC default (docker/Dockerfile, docker/Dockerfile.studio, docker/entrypoint.sh):
  - Make cu12.8 the immutable baked default (libnvrtc.so.12 -> .cu128.orig) with
    a staged .cu13 alias, and have select_cuda_jit_tools retarget to cu13 ONLY
    for sm_103/sm_121. Previously cu13 was baked as the default and restored to
    cu12.8 at runtime, so a non-root `docker run --user` container that cannot
    rewrite the symlink stayed on cu13 NVRTC and emitted cubins a 570-579 driver
    cannot load. The safe default now needs no runtime write.

Adds regression tests for each case (tests/python/test_unsloth_pip_shim.py,
tests/sh/test_select_cuda_jit_tools.sh).
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Chef's kiss.

Reviewed commit: 6a078b1a45

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants